// all trax programs must include trax.hpp
#include "trax.hpp"
#include "Image.h"
#include "Vector3.h"
#include "Sphere.h"
#include "HitRecord.h"
#include "PinHoleCamera.h"
#include "Scene.h"
#include "Shader.h"
#include "AABBox.h"
#include "Unit_Test.h"
#include "atTriangle.h"
#include "RayTracer.h"
#include "atUtil.h"
#include "atMaterial.h"
#include "PhotonMap.h"
#include "atAtomicReg.h"
#include "HitMap.h"
#include "atMathConstant.h"

#define MAX_MSAA 128
#define PHT_RADIUS 1.0f
#define ALPHA_R 1.0f
#define PHT_PASS 1

#define ONLY_PHOTON_MAP 1 // only compute photon map, no rendering
#define RUN_TEST 0
#define PRINT_PERCENTAGE 0

#define ADD_DIRECT_ILLUM 1

// "main" function is called "trax_main", returns void instead of int
void trax_main()
{
    using namespace AtomTrace;
#if RUN_TEST
    printf("Unit test\n");
	//Unit_test_Vector3();
    //Unit_test_BoxIntersection();
    //Unit_test_TriangleIntersection();
    //Unit_test_BVH_Stack();
    //Unit_test_GetMemInfo();
    //Unit_test_Atomicadd();
    //Unit_test_PosToZcurve();
    Unit_test_RegularPhotonMap();
    return;
#endif
    
    const int threadId = atomicinc(REG_THREAD_ID);

    const int maxDepth = 7;

    float color_threshold = 0.001f;
	int xres = GetXRes();
	int yres = GetYRes();
	int start_fb = GetFrameBuffer();
    int num_samples = loadi(TRAX_NUM_SAMPLES);
    if(num_samples > MAX_MSAA)
    {
        printf("Warning, maximum number of samples for MSAA is %d\n", MAX_MSAA);
        num_samples = MAX_MSAA;
    }
    float num_samplesF_inv = 1.0f / static_cast<float>(num_samples);
	Image image(xres, yres, start_fb);
    HitMap hitmap;
	

    //<Scene def. lab3>------------------
    PinHoleCamera pcamera;
    PointLight light;
    //Vector3 lpos;
    Scene scene;
    atTriangle tempTri;

    scene.numLights = 1;
    //light..LoadFromGM(GetLight());
    //scene.pointLights = &light;
    scene.startLightAddr = GetLight();
    //lpos.LoadFromGM(GetLight());
    scene.startMatAddr = GetMaterials();
    pcamera.LoadFromGM(GetCamera());
    scene.startTriAddr = GetStartTriangles();
    scene.numTri = GetNumTriangles();
    scene.bvh_start = GetBVH();
    //-----------------------</Scene def. lab3>
    
    PhotonMap pm;
    pm.SetMaterialStartAddr(scene.startMatAddr);
    

    //light.LoadFromGM(scene.startLightAddr);
    //printf("lpos: [%f, %f, %f]\n", light.pos[0], light.pos[1], light.pos[2]);

    int imgSize = xres * yres;
    int subImg = 1;
    int bImgSize = imgSize/subImg;
    //Vector3 des;
    //int nodeID = 0;
    int percent = 0;
    Color res, red, sumres;
    Color tresult[MAX_MSAA];
    //red.Set(1.0f, 0.0f, 0.0f);

    // ---------- Loop through photon pass --------------------
    for (int ppass=0; ppass < PHT_PASS; ppass++) {
        
        //<PhotonMap test> ------------------------------
        pm.Compute(scene, threadId);
        if(threadId == 0){
            atomicinc(REG_PHT_PASS);
        }
        barrier(REG_BARRIER);
#if ONLY_PHOTON_MAP
        //pm.Reset();
    	if(threadId==0)
    		printf("--------------- No rendering!! ----------------\n");
        pm.Sort();
        return;
#endif
        
        //--------------------------------</PhotonMap test>*/
        if(threadId == 0)
            printf("\nthread %d, photon pass %d/%d\n", threadId, ppass+1, PHT_PASS);
    
        // ---------- Loop through every pixels in z-shape pattern ---------
        for(int pixww = atomicinc(REG_PIXEL_ID)*subImg;
                pixww < imgSize;
                pixww = atomicinc(REG_PIXEL_ID) * subImg)
        //int pixww = 63; //256+512 * 256;
        {
#if PRINT_PERCENTAGE
            float pecf = static_cast<float>(pixww) * 100.0f / static_cast<float>(imgSize);
            if(pixww % 4 == 0)
                printf("\rphoton indirrect %.2f", pecf);
#endif
            for(int subPix = 0; subPix<subImg; subPix++)
            {
                int pix = pixww+subPix;
                int ix, iy;
                ConvertLineIndexToZindexXY(pix, ix, iy);
                
                float offx = 0.0f;
                float offy = 0.0f;
                Ray ray;
                int outpix;
                outpix = iy * xres + ix;
                res.Set(0.0f, 0.0f, 0.0f);

                // ---- multi sampling anti-aliasing ---------
                int num_used_samples = num_samples;
                Color average;
                // TODO: progressive photon map
                float maxR2;
                
                HitPoint curHit;
                hitmap.Get(outpix, curHit);
                HitPoint newHp;
                if(ppass == 0)
                {
                    curHit.r = PHT_RADIUS; // Initial estimate radius
                    newHp.r = PHT_RADIUS;
                    
                    curHit.flux.Set(0.0f);
                }
                
                newHp.flux.Set(0.0f);
                curHit.pix = outpix;
                newHp.pix = outpix;
                //for(int i=0; i<num_samples; i++)
                {
                    int i=0;
                    //offx = (trax_rand() - 0.5f);
                    //offy = (trax_rand() - 0.5f);
                    
                    // Halton sequence offset
                    offx = Halton(i, 2);
                    offy = Halton(i, 3);

                    // Center them
                    offx += (offx>0.5f)? -1.0f : 0.0f;
                    offy += (offy>0.5f)? -1.0f : 0.0f;

                    ray = pcamera.CastRay(ix+offx, iy+offy);
                    int depth = 0;

                    HitRecord hit;
                    
                    //tresult[i].Set(0.0f, 0.0f, 0.0f);
                    while(depth == 0 || (depth < maxDepth && hit.newDir.LengthSquare() > 0.1f))
                    {
                        depth++;
                        hit.Init();
                        bool isHit = false;
                        isHit = RayTracer::TraceBVH(ray, hit);
                        if(isHit)
                        {
                            //tresult[i].Set(1.0f, 1.0f, 1.0f);
                            //tresult[i] += shade(ray, hit, scene);
                            
                            //Color pht_flux;
                            pm.RadianceEstimate(hit, curHit.r, newHp.flux, maxR2, newHp.npht);
                            newHp = hit;
                            newHp.r = sqrt(maxR2);
                            if(hit.matIllum != 3 && hit.matIllum != 6 ){
                                //pht_flux *= 1.0f / (3.141592653f * maxR2);
                                //tresult[i] += pht_flux;
                                if(ppass == 0)
                                {
                                    //curHit = newHp;
                                }
                                //else
                                {
                                    float Npht = curHit.npht;
                                    float Mpht = newHp.npht;
                                    float facf = (Npht + ALPHA_R * Mpht) / (Npht + Mpht);
                                    curHit.r *= sqrt(facf);
                                    curHit.npht += ALPHA_R * Mpht;
                                    curHit.flux = (curHit.flux + newHp.flux) * facf;
                                }
                            }
                            
                            Material::GetNewDir(hit.matIllum, hit, hit.newDir, false);

                            ray._o = hit.P + hit.newDir.GetNormalize() * 0.01f;
                            ray._dir = hit.newDir;
                        }else
                        {
                            break;
                        }
                        //tresult[i] *= 0.8f;
                    } // bounce rays
                    
                    //tresult[i].Clamp(0.0f, 1.0f);
                    //res += tresult[i];
                    hitmap.Set(outpix, curHit);
                    res += curHit.flux;
                    
                    // Adaptive multi-sampling
                    average = res * (1.0f / static_cast<float>(i+1));
                    Color variance;
                    variance.Set(0.0f, 0.0f, 0.0f);
                    /**
                    if(i>=4){
                        for(int j=0; j<i; j++)
                        {
                            Color tempDiff = tresult[j] - average;
                            variance += tempDiff * tempDiff;
                        }
                        variance *= (1.0f / static_cast<float>(i+1));
                        if((variance.r() < color_threshold && 
                            variance.g() < color_threshold &&
                            variance.b() < color_threshold) || 
                            i == num_samples-1)
                        {
                            //res = average;
                            num_used_samples = i+1;

                            break;

                        }
                    } // anti aliasing estimation */
                } // loop through number of samples, anti-aliasing

                //image.set(outpix, average);
            } // loop through pixels in block
        } // loop through pixels block
        
        barrier(REG_BARRIER);
        if(threadId == 0)
        {
            int pixId = atomicadd(REG_PIXEL_ID, 0);
            pixId = atomicadd(REG_PIXEL_ID, -pixId);
            printf("\nReset register for pixel_ID: %d, for next photon pass\n", pixId);
        }
        pm.Reset();
        barrier(REG_BARRIER);
    } // loop through photon pass
    
    /**
    if(threadId == 0)
    {
        int pixId = atomicadd(REG_PIXEL_ID, 0);
        pixId = atomicadd(REG_PIXEL_ID, -pixId);
        printf("\nReset register for pixel_ID: %d, for next photon pass\n", pixId);
    }
    barrier(REG_BARRIER);
     // */

#if ADD_DIRECT_ILLUM
    // Direct illuminations
    // ---------- Loop through every pixels in z-shape pattern ---------
    for(int pixww = atomicinc(REG_PIXEL_ID)*subImg;
        pixww < imgSize;
        pixww = atomicinc(REG_PIXEL_ID) * subImg)
        //int pixww = 63; //256+512 * 256;
    {
#if PRINT_PERCENTAGE
        float pecf = static_cast<float>(pixww) * 100.0f / static_cast<float>(imgSize);
        if(pixww % 4 == 0)
            printf("\rrendering %.2f", pecf);
#endif
        for(int subPix = 0; subPix<subImg; subPix++)
        {
            int pix = pixww+subPix;
            int ix, iy;
            ConvertLineIndexToZindexXY(pix, ix, iy);
            
            float offx = 0.0f;
            float offy = 0.0f;
            Ray ray;
            int outpix;
            outpix = iy * xres + ix;
            res.Set(0.0f, 0.0f, 0.0f);
            
            // ---- multi sampling anti-aliasing ---------
            int num_used_samples = num_samples;
            Color average;
            
            for(int i=0; i<num_samples; i++)
            {
                //int i=0;
                //offx = (trax_rand() - 0.5f);
                //offy = (trax_rand() - 0.5f);
                
                // Halton sequence offset
                offx = Halton(i, 2);
                offy = Halton(i, 3);
                
                // Center them
                offx += (offx>0.5f)? -1.0f : 0.0f;
                offy += (offy>0.5f)? -1.0f : 0.0f;
                
                ray = pcamera.CastRay(ix+offx, iy+offy);
                int depth = 0;
                
                HitRecord hit;
                
                tresult[i].Set(0.0f, 0.0f, 0.0f);
                while(depth == 0 || (depth < maxDepth && hit.newDir.LengthSquare() > 0.1f))
                {
                    depth++;
                    hit.Init();
                    bool isHit = false;
                    isHit = RayTracer::TraceBVH(ray, hit);
                    if(isHit)
                    {
                        //tresult[i].Set(1.0f, 1.0f, 1.0f);
                        tresult[i] += shade(ray, hit, scene);
                        
                        Material::GetNewDir(hit.matIllum, hit, hit.newDir, false);
                        
                        ray._o = hit.P + hit.newDir.GetNormalize() * 0.01f;
                        ray._dir = hit.newDir;
                    }else
                    {
                        tresult[i] += scene.background;
                        break;
                    }
                    //tresult[i] *= 0.8f;
                } // bounce rays
                
                tresult[i].Clamp(0.0f, 1.0f);
                res += tresult[i];
                
                // Adaptive multi-sampling
                average = res * (1.0f / static_cast<float>(i+1));
                Color variance;
                variance.Set(0.0f, 0.0f, 0.0f);
                /**/
                 if(i>=4){
                     for(int j=0; j<i; j++)
                     {
                         Color tempDiff = tresult[j] - average;
                         variance += tempDiff * tempDiff;
                     }
                     variance *= (1.0f / static_cast<float>(i+1));
                     if((variance.r() < color_threshold &&
                         variance.g() < color_threshold &&
                         variance.b() < color_threshold) ||
                         i == num_samples-1)
                     {
                         //res = average;
                         num_used_samples = i+1;
                     
                         break;
                     
                     }
                 } // anti aliasing estimation */
            } // loop through number of samples, anti-aliasing
            
            image.set(outpix, average);
        } // loop through pixels in block
    } // loop through pixels block
    
#endif // if ADD_DIRECT_ILLUM
    
    barrier(REG_BARRIER);
    if(threadId == 0)
    {
        int pixId = atomicadd(REG_PIXEL_ID, 0);
        pixId = atomicadd(REG_PIXEL_ID, -pixId);
        printf("\nReset register for pixel_ID: %d for final output\n", pixId);
    }
    barrier(REG_BARRIER);

    
    Color d_illm;
    // Copy hit map result to frame buffer
    for (int pix = atomicinc(REG_PIXEL_ID); pix<imgSize; pix = atomicinc(REG_PIXEL_ID))
    {
        HitPoint hit;
        hitmap.Get(pix, hit);
        d_illm = image.get(pix);
        if(PHT_PASS > 0){
            float r2 = hit.r * hit.r;
            if (r2 > 0.01f) {
                hit.flux /= (PI_F * r2 * static_cast<float>(PHT_PASS) * GetSingleEmittedPhotons() );
            }else{
                hit.flux = 0.0f;
            }
            
        }
        //hit.flux *= 4.0f * PI_F;
        image.set(pix, hit.flux+d_illm);
    } // Copy hit map result to frame buffer */
}


